#!/usr/bin/env python
# -*- coding:utf-8 -*-
from unicorn.x86_const import *
from unicorn import *
from capstone import *
import struct
import time
from lib.utils import printer

def tools_disassembly(uc, addr, size):
    CODE = uc.mem_read(addr, size)
    md = Cs(CS_ARCH_X86, CS_MODE_32)
    for i in md.disasm(bytes(CODE), addr):
        printer.less("%x:%s  %s" %(i.address, i.mnemonic, i.op_str))
    print("")
def get_win32_api_by_addr(win32Dict:dict,addr):
    for _,dll in win32Dict.items():
        if dll['dllBase'] <= addr <= dll['dllLimit']:
            apiName = dll['apiDict'].get(addr)
            return apiName
    return None
def is_valid_zone(codeZone,dllZone,stackZone,address):
    if codeZone['base'] <= address <= codeZone['limit'] or \
            stackZone['base']<= address <= stackZone['limit'] or \
            dllZone['base'] <= address <= dllZone['limit']:
            return True
    return False
def pe_hook_code(uc, address, size, emu):
    sysBase = 0xff00000
    printer.info('>>> 指令执行地址 0x%x <<< | >>> 指令大小:0x%x <<<' % (address, size),flag='')
    esp = uc.reg_read(UC_X86_REG_ESP)
    args = struct.unpack("<IIIIII",uc.mem_read(esp,24))
    callAddr = args[0] - 6
    printer.test("-------------------- 寄存器情况 --------------------",flag='')
    emu.reg.print_registers()
    if sysBase <= address <= sysBase + sysBase * 0x100:
        apiName = emu.get_win32_api_by_addr(address)
        if apiName == None:
            print('[!] %x: 没有分配的API地址 @ %x' % (callAddr, address))
        else:
            print('\n[+] 函数地址: %x: 调用的API函数 %s' % (callAddr, apiName))
            print('[+] -------------------- 参数压栈情况 --------------------')
            for i in range(1,5):
                if emu.valid_zone(args[i]):
                    argstr = uc.mem_read(args[i],30).decode('utf8', errors='ignore').strip('\x00')
                else:
                    argstr = "0"
                print('>>> args[参数]_%i( [栈址] %x) --> [内存地址] %.8x | [参数值] %s' % (i, esp + 4 * i, args[i], argstr))
            print("")
    else:
        tools_disassembly(uc,address,size)
    time.sleep(1)
def hook_mem_invalid(uc, access, address, size, value, user_data):
    eip = uc.reg_read(UC_X86_REG_EIP)
    if access == UC_MEM_WRITE:
        printer.warn("无效的内存写入[UC_MEM_WRITE]! 内存地址: 0x%x, EIP = 0x%X, 数据大小 = %u, 数据值 = 0x%x" % (address, eip, size, value))
    if access == UC_MEM_READ:
        printer.warn("无效的内存读取[UC_MEM_READ]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_FETCH:
        printer.warn("无效的内存执行事件[UC_MEM_FETCH]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_READ_UNMAPPED:
        printer.warn("无效的内存读取[UC_MEM_READ_UNMAPPED] 内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_WRITE_UNMAPPED:
        printer.warn("无效的内存写入[UC_MEM_WRITE_UNMAPPED] 内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_FETCH_UNMAPPED:
        printer.warn("无效的内存执行事件[UC_MEM_FETCH]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_WRITE_PROT:
        printer.warn("内存写保护写入事件无效[UC_MEM_WRITE_PROT]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_HOOK_MEM_READ_PROT:
        printer.warn("内存读保护读取事件无效[UC_HOOK_MEM_READ_PROT]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_FETCH_PROT:
        printer.warn("内存保护执行事件无效[UC_MEM_FETCH_PROT]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    if access == UC_MEM_READ_AFTER:
        printer.warn("内存读取出发回调[UC_MEM_READ_AFTER]！内存地址: 0x%x; EIP = 0x%X, 数据大小 = %u" % (address, eip, size))
    return False

def code_hook_code(uc, address, size, emu):
    sysBase = 0xff00000
    print('>>> 指令执行地址 0x%x <<< | >>> 指令大小:0x%x <<<' % (address, size))
    print("-------------------- 寄存器情况 --------------------")
    emu.reg.print_registers()
    esp = uc.reg_read(UC_X86_REG_ESP)
    args = struct.unpack("<IIIIII",uc.mem_read(esp,24))
    callAddr = args[0] - 6
    print('[+] -------------------- 参数压栈情况 --------------------')
    for i in range(1,5):
        if emu.valid_zone(args[i]):
            argstr = uc.mem_read(args[i],30).decode('utf8', errors='ignore').strip('\x00')
        else:
            argstr = "0"
        print('>>> args[参数]_%i( [栈址] %x) --> [内存地址] %.8x | [参数值] %s' % (i, esp + 4 * i, args[i], argstr))
    if sysBase <= address <= sysBase + sysBase * 0x100:
        apiName = emu.get_win32_api_by_addr(address)
        if apiName == None:
            print('[!] %x: 没有分配的API地址 @ %x' % (callAddr, address))
        else:
            emu.emu_stop()
            print('\n[+] 函数地址: %x: 调用的API函数 %s' % (callAddr, apiName))
    tools_disassembly(uc,address,size)
    time.sleep(1)

def code_hook_code_test(uc, address, size, emu):
    sysBase = 0xff00000
    print('>>> 指令执行地址 0x%x <<< | >>> 指令大小:0x%x <<<' % (address, size))
    print("-------------------- 寄存器情况 --------------------")
    emu.reg.print_registers()
    esp = uc.reg_read(UC_X86_REG_ESP)
    args = struct.unpack("<IIIIII",uc.mem_read(esp,24))
    callAddr = args[0] - 6
    print(hex(callAddr))
    print('[+] -------------------- 参数压栈情况 --------------------')
    for i in range(0,5):
        try:
            argstr = uc.mem_read(args[i],30).decode('utf8', errors='ignore').strip('\x00')
            if not argstr:
                argstr = "0"
        except Exception as e:
            argstr = "异常"
        print('>>> args[参数]_%i( [栈址] %x) --> [内存地址] %.8x | [参数值] %s' % (i, esp + 4 * i, args[i], argstr))
    tools_disassembly(uc,address,size)
    time.sleep(1)
def code_hook_interrupt(uc,intno,user_data):
    if intno == 0x80:
        eax = uc.reg_read(UC_X86_REG_EAX)
        if eax == 11:
            printer.warn("接收到中断指令,根据EAX寄存器的值判断调用的功能为sys_execve功能即运行指定程序!")
    